package com.pkk.cloud.support.gateway.limiter.global.util;

import static com.pkk.cloud.support.gateway.limiter.global.constand.GlobalLimitConstand.DEFAULT_GLOBAL_METHOD_LIMIT;

import com.google.common.util.concurrent.RateLimiter;
import com.pkk.cloud.support.gateway.limiter.global.properties.GlobalRateLimitProperties;
import java.util.concurrent.ConcurrentHashMap;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.CommandLineRunner;

/**
 * @description: 接口的限制工具类
 * @author: peikunkun
 * @create: 2020-03-24 15:18
 **/
@Slf4j
public class GlobalMethodRateLimitUtil implements CommandLineRunner {

  @Autowired
  private GlobalRateLimitProperties globalRateLimitProperties;
  //限制url接口的map
  private final ConcurrentHashMap<String, RateLimiter> map = new ConcurrentHashMap<>();

  /**
   * 接口的限制
   */
  private static Integer METHOD_LIMIT = -1;


  /**
   * @Description: 初始化操作
   * @Param:
   * @return: void
   * @Author: peikunkun
   * @Date: 2020/3/24 0024 下午 4:59
   */
  public void init() {
    if (null != globalRateLimitProperties && null != globalRateLimitProperties.getMethod() &&
        globalRateLimitProperties.getMethod().isEnable()) {
      METHOD_LIMIT = globalRateLimitProperties.getMethod().getLimitNum();
    }

    if (null == METHOD_LIMIT) {
      METHOD_LIMIT = DEFAULT_GLOBAL_METHOD_LIMIT;
    }

    //小于零设置理想型的无限大
    if (METHOD_LIMIT < 0) {
      METHOD_LIMIT = 999999;
    }
    log.info("[网关服务]-[全局每秒请求接口最大限制次数]-[{}]", METHOD_LIMIT);
  }

  /**
   * @Description: 尝试获取调用接口的是否可调用
   * @Param: functionName
   * @return: boolean
   * @Author: peikunkun
   * @Date: 2020/3/24 0024 下午 3:26
   */
  public boolean tryAcquire(String functionName) {
    if (METHOD_LIMIT < 0) {
      return true;
    }

    //获取rateLimiter
    RateLimiter rateLimiter = null;
    if (map.containsKey(functionName)) {
      rateLimiter = map.get(functionName);
    } else {
      map.put(functionName, RateLimiter.create(METHOD_LIMIT));
      rateLimiter = map.get(functionName);
    }
    return rateLimiter.tryAcquire();
  }


  @Override
  public void run(String... args) throws Exception {
    //初始化配置
    this.init();
  }
}
